# -*- coding: utf-8 -*-

# File Name : corner_t_bbox_generator
# Description :
# Author : marje
# date : 2022/6/5
# Change Activity : 2022/6/5
import cv2
import numpy as np

from ml.cv.algorithms.bbox_generate.corner_to_bbox.bbox_correct import BboxCorrector
from ml.cv.algorithms.sort_points.sort_points import resort_points


class CornerToBboxGenerator:
    """
    Mask to bbox.
    """

    def __init__(self, text_area_pixel, threshold=0.4, min_pixel_size=50, enlarge_ratio=2, enlarge_pixel=2):
        self.text_area_pixel = text_area_pixel
        self.threshold = threshold
        self.enlarge_ratio = enlarge_ratio
        self.min_pixel_size = min_pixel_size
        self.format_pixel = self._format_pixel()
        self.enlarge_pixel = enlarge_pixel
        self.connected_domain = self.__get_connected_domains()

    def _get_bbox_from_one_domain(self, domain_value):
        """

        :param domain_value:
        :return:
        """
        locations = np.where(self.connected_domain == domain_value)
        if locations[0].size < self.min_pixel_size:
            return None
        bg = np.zeros_like(self.connected_domain, dtype=np.uint8)
        bg[locations] = 255
        return self._gen_one_bbox_from_single_mask(bg)

    @staticmethod
    def _enlarge_bbox(enlarge_pixel, bbox):
        """
        enlarge bbox
        :param enlarge_pixel:
        :param bbox:
        :return:
        """
        p1, p2, p3, p4 = bbox[0], bbox[1], bbox[2], bbox[3]
        p1 = [p1[0] - enlarge_pixel, p1[1] - enlarge_pixel]
        p2 = [p2[0] + enlarge_pixel, p2[1] - enlarge_pixel]
        p3 = [p3[0] + enlarge_pixel, p3[1] + enlarge_pixel]
        p4 = [p4[0] - enlarge_pixel, p4[1] + enlarge_pixel]
        bbox = [p1, p2, p3, p4]
        for point in bbox:
            width, height = point
            if width < 0:
                point[0] = 0
            if height < 0:
                point[1] = 0
        new_bbox = np.array([p1, p2, p3, p4])
        return new_bbox

    def _gen_one_bbox_from_single_mask(self, bg):
        bbox_corrector = BboxCorrector(bg)
        bbox = bbox_corrector.get_box()
        if bbox is None:
            return None
        bbox = bbox.astype(np.int) * self.enlarge_ratio
        bbox = self._enlarge_bbox(self.enlarge_pixel, bbox)
        return bbox

    def gen_bboxes(self):
        domain_max_value = self.connected_domain.max()
        if domain_max_value == 0:
            return []
        bboxes = []
        for domain_value in range(1, domain_max_value + 1):  # for each domain
            box = self._get_bbox_from_one_domain(domain_value)
            if box is None:
                continue
            box = resort_points(box)
            bboxes.append(box)
        return bboxes

    def _format_pixel(self):
        bg = np.zeros_like(self.text_area_pixel, dtype=np.uint8)
        bg[np.where(self.text_area_pixel > self.threshold)] = 255
        return bg

    def __get_connected_domains(self):
        """
        combine and return connected domains
        :return:
        """
        _, labels_image = cv2.connectedComponents(self.format_pixel, connectivity=4)
        return labels_image
